import json
import time

from flask import Flask, Response
import cv2
import threading
import socket
import struct
import numpy as np

from logger import log

app = Flask(__name__)

# 存储视频帧
"""
{
    "vs1": {
        "client_socket": client_socket
        "streams": {
            id:{
                "name": "name",
                "buf": buf,
                "count": 0,
            }
        },
        "last_heartbeat": time.time()
    },
}
"""
edges = {}
HEARTBEAT_INTERVAL = 30
HEARTBEAT_TIMEOUT = 60


def get_frame(edge_code, id):
    while True:
        frame = edges[edge_code]['streams'][id]['frame']
        if frame is not None:
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
        else:
            yield b''
        time.sleep(0.2)


@app.route('/')
def hello():
    links = []
    for edge_key in edges.keys():
        streams = edges[edge_key]['streams']
        for key in streams:
            links.append(f'<li><a href="/{edge_key}/{key}">{streams[key]["name"]}</a></li>')
    links_html = '<ul>' + ''.join(links) + '</ul>'
    return f'<html><body><h1>Streams</h1>{links_html}</body></html>'


@app.route('/<edge_code>/<id_str>')
def video_feed(edge_code, id_str):
    # 客户端连接，通知推流终端开始推流
    log.info(f"video_feed:{edge_code},{id_str}")
    id = int(id_str)
    if edge_code not in edges:
        log.info(f"没有找到客户端:{edge_code}")
        return
    client_socket = edges[edge_code]['client_socket']
    stream = edges[edge_code]['streams'][id]
    stream['count'] += 1
    if stream['count'] == 1:
        log.info(f"发送连接请求:{id},stream['count']:{stream['count']}")
        client_socket.sendall(json.dumps({"cmd": "start", "id": id}).encode())

    def generate():
        try:
            for frame in get_frame(edge_code, id):
                yield frame
        except GeneratorExit:
            log.info(f"GeneratorExit triggered for stream {id}")
        finally:
            # 结束推流
            stream['count'] -= 1
            log.info(f"客户端退出:{id},stream['count']:{stream['count']}")
            if stream['count'] == 0:
                log.info(f"发送关闭连接请求:{id}")
                client_socket.sendall(json.dumps({"cmd": "stop", "id": id}).encode())

    response = Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

    @response.call_on_close
    def close_generator():
        log.info(f"Closing generator for stream {id}")
        generate().close()

    return response


def handle_client_connection(client_socket, edge_code):
    log.info(f"handle_client_connection:{edge_code}")
    payload_size = struct.calcsize("QQ")

    def heartbeat_check():
        while True:
            if time.time() - edges[edge_code]['last_heartbeat'] > HEARTBEAT_TIMEOUT:
                log.warning(f"Heartbeat timeout for {edge_code}, disconnecting...")
                client_socket.close()
                del edges[edge_code]
                break
            time.sleep(HEARTBEAT_INTERVAL)

    heartbeat_thread = threading.Thread(target=heartbeat_check, daemon=True)
    heartbeat_thread.start()

    while True:
        try:
            # 读取包头
            delimiter = client_socket.recv(4)
            if not delimiter or delimiter != b'@#@#':
                continue
            # 读取帧头
            header = client_socket.recv(payload_size)
            if len(header) < payload_size:
                continue

            id, frame_size = struct.unpack("QQ", header)
            if id == 0 and frame_size == 0:
                # 心跳包
                edges[edge_code]['last_heartbeat'] = time.time()
                client_socket.sendall(json.dumps({"cmd": "heartbeat"}).encode())
                continue

            # log.info(f"Receiving frame for id {id}, size {frame_size}")

            # 读取帧数据
            frame_data = b""
            while len(frame_data) < frame_size:
                packet = client_socket.recv(frame_size - len(frame_data))
                if not packet:
                    break
                frame_data += packet

            if len(frame_data) != frame_size:
                log.info("Incomplete frame received")
                continue
            frame = np.frombuffer(frame_data, dtype=np.uint8)
            frame = cv2.imdecode(frame, cv2.IMREAD_COLOR)
            edges[edge_code]['streams'][id]['frame'] = frame
        except Exception as e:
            log.exception(e)
            break

    log.info(f"Connection closed for {edge_code}")
    client_socket.close()
    if edge_code in edges:
        del edges[edge_code]


def start_tcp_server():
    while True:
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            server_socket.bind(('0.0.0.0', 9999))
            server_socket.listen(100)
            log.info("TCP server listening on port 9999")
            while True:
                client_socket, addr = server_socket.accept()
                log.info(f"Accepted connection from {addr}")
                try:
                    # 读取4字节数据长度
                    length_data = client_socket.recv(4)
                    if len(length_data) != 4:
                        log.error("Invalid length header")
                        client_socket.close()
                        continue
                    data_length = struct.unpack('!I', length_data)[0]

                    # 读取完整JSON数据
                    data = b''
                    while len(data) < data_length:
                        chunk = client_socket.recv(data_length - len(data))
                        if not chunk:
                            break
                        data += chunk
                    if len(data) != data_length:
                        log.error("Incomplete data received")
                        client_socket.close()
                        continue

                    # 解析JSON
                    obj = json.loads(data.decode())
                    edge_code = obj['edge_code']
                    streams = {}
                    for stream in obj['streams']:
                        streams[stream['id']] = {"name": stream['name'], "frame": None, "count": 0}
                    edges[edge_code] = {"client_socket": client_socket, "streams": streams,
                                        "last_heartbeat": time.time()}

                    # 启动客户端处理线程
                    client_handler = threading.Thread(
                        target=handle_client_connection,
                        args=(client_socket, edge_code), daemon=True
                    )
                    client_handler.start()
                except Exception as e:
                    log.exception(e)
                    client_socket.close()
        except Exception as e:
            log.exception(e)
            log.info("Waiting 10 seconds before retrying...")
            time.sleep(10)


if __name__ == '__main__':
    tcp_thread = threading.Thread(target=start_tcp_server, daemon=True)
    tcp_thread.start()

    app.run(host='0.0.0.0', port=8888)
